import pandas as pd
import numpy as np
from MyTT import *
from MyTT_Patch import *
import torch

class TradeSignalCalculator:
    def __init__(self, data):
        """
        初始化
        :param data: 输入的DataFrame数据
        """
        self.data = data.copy()
    
    def three_stars_bottom(self):
        """
        三星底部形态识别因子
        计算STARK信号基于以下条件：
        - A3: abs(close(t-2) - open(t-2)) / open(t-2) < 0.012
        - A2: (open(t-1) - close(t-1)) / open(t-1) < 0.008
        - A1: abs(close - open) / open < 0.012
        - B1: close(t-2) > open(t-2) ? open(t-2) : close(t-2)
        - B2: close(t-1) < open(t-1) ? open(t-1) : close(t-1)
        - B3: close > open ? open : close
        - B4: B1 > B2 AND B3 > B2 AND close(t-1) < MA(close, 5)
        - STARK: A3 AND A2 AND A1 AND B4
        """
        if len(self.data) < 5:
            print(f"Group too small ({len(self.data)} rows); skipping three_stars_bottom.")
            self.data['STARK'] = False
            return self
        
        required_cols = ['open', 'close']
        missing_cols = [col for col in required_cols if col not in self.data.columns]
        if missing_cols:
            print(f"Missing columns {missing_cols}; skipping three_stars_bottom.")
            self.data['STARK'] = False
            return self
        
        try:
            ref_open_2 = self.data['open'].shift(2)
            ref_close_2 = self.data['close'].shift(2)
            ref_open_1 = self.data['open'].shift(1)
            ref_close_1 = self.data['close'].shift(1)
            open_current = self.data['open']
            close_current = self.data['close']
            
            self.data['three_stars_a3'] = np.where(
                ref_open_2 != 0,
                (abs(ref_close_2 - ref_open_2) / ref_open_2) < 0.012,
                False
            )
            
            self.data['three_stars_a2'] = np.where(
                ref_open_1 != 0,
                (ref_open_1 - ref_close_1) / ref_open_1 < 0.008,
                False
            )
            
            self.data['three_stars_a1'] = np.where(
                open_current != 0,
                (abs(close_current - open_current) / open_current) < 0.012,
                False
            )
            
            self.data['three_stars_b1'] = np.where(
                ref_close_2 > ref_open_2,
                ref_open_2,
                ref_close_2
            )
            
            self.data['three_stars_b2'] = np.where(
                ref_close_1 < ref_open_1,
                ref_open_1,
                ref_close_1
            )
            
            self.data['three_stars_b3'] = np.where(
                close_current > open_current,
                open_current,
                close_current
            )
            
            ma5 = pd.Series(MA(self.data['close'], 5), index=self.data.index)
            self.data['three_stars_b4'] = (
                (self.data['three_stars_b1'] > self.data['three_stars_b2']) &
                (self.data['three_stars_b3'] > self.data['three_stars_b2']) &
                (ref_close_1 < ma5.shift(1))
            )
            
            self.data['STARK'] = (
                self.data['three_stars_a3'] &
                self.data['three_stars_a2'] &
                self.data['three_stars_a1'] &
                self.data['three_stars_b4']
            )
            
            return self
            
        except Exception as e:
            print(f"Error in three_stars_bottom: {e}")
            self.data['STARK'] = False
            return self
    
    def buy_potential(self):
        if len(self.data) < 34: # For HHV/LLV 27 and 34
            print(f"Group too small ({len(self.data)} rows) for buy_potential; skipping.")
            self.data['POTENTIALBUY'] = False
            return self
        
        required_cols = ['low', 'high', 'close']
        missing_cols = [col for col in required_cols if col not in self.data.columns]
        if missing_cols:
            print(f"Missing columns {missing_cols}; skipping buy_potential.")
            self.data['POTENTIALBUY'] = False
            return self
        
        # Ensure LLV and HHV return Series to use .replace/.fillna
        # Or handle them as numpy arrays
        self.data['VAR05'] = LLV(self.data['low'], 27)
        self.data['VAR06'] = HHV(self.data['high'], 34)
        
        # Avoid division by zero
        denominator = self.data['VAR06'] - self.data['VAR05']
        # --- START MODIFICATION ---
        # 使用 np.where 代替 .replace 和 .fillna 处理 numpy 数组
        denominator = np.where(denominator == 0, 0.0001, denominator)
        # --- END MODIFICATION ---
        
        self.data['VAR07'] = EMA((self.data['close']-self.data['VAR05'])/denominator*4,4)*25
        self.data['建仓区'] = self.data['VAR07']<10
        self.data['建仓区连续天数'] = self.data['建仓区'].astype(int)
        self.data['建仓区连续天数'] = self.data['建仓区连续天数'] * (self.data['建仓区连续天数'].groupby((self.data['建仓区连续天数'] != self.data['建仓区连续天数'].shift()).cumsum()).cumcount() + 1)
        self.data['建仓区满足3天'] = self.data['建仓区连续天数'] >= 3

        # For VAR1, LLV/HHV 5
        if len(self.data) < 5:
            print(f"Group too small ({len(self.data)} rows) for VAR1 in buy_potential; skipping.")
            self.data['POTENTIALBUY'] = False
            return self
            
        llv_low_5 = LLV(self.data['low'], 5)
        hhv_high_5 = HHV(self.data['high'], 5)
        
        # Avoid division by zero
        denominator_var1 = hhv_high_5 - llv_low_5
        # --- START MODIFICATION ---
        # 使用 np.where 代替 .replace 和 .fillna 处理 numpy 数组
        denominator_var1 = np.where(denominator_var1 == 0, 0.0001, denominator_var1)
        # --- END MODIFICATION ---

        val = (self.data['close']-llv_low_5)/denominator_var1*100
        sma_val_5_1 = SMA(val, 5, 1)
        sma_sma_val_3_1 = SMA(sma_val_5_1, 3.2, 1) # Note: SMA(X, N, M) is usually (M*X + (N-M)*REF(X,1))/N. Here M=1.
        
        self.data['VAR1'] = 4*sma_val_5_1 - 3*sma_sma_val_3_1
        
        self.data['上升'] = CROSS_TORCH(self.data['VAR1'], 8)
        self.data['POTENTIALBUY'] = (self.data['建仓区满足3天'] | self.data['建仓区满足3天'].shift(1)) & self.data['上升']
        return self
    
    def buy_entry_signal(self):
        if len(self.data) < 10: # For EMA 10
            print(f"Group too small ({len(self.data)} rows) for buy_entry_signal; skipping.")
            self.data['ENTRYBUY'] = False
            return self

        required_cols = ['close']
        missing_cols = [col for col in required_cols if col not in self.data.columns]
        if missing_cols:
            print(f"Missing columns {missing_cols}; skipping buy_entry_signal.")
            self.data['ENTRYBUY'] = False
            return self
        
        self.data['ema_diff_1'] = EMA(self.data['close'], 2) - EMA(self.data['close'], 10)
        self.data['ema_diff_2'] = EMA(self.data['ema_diff_1'], 2)        
        self.data['ENTRYBUY'] = CROSS_TORCH(self.data['ema_diff_1'], self.data['ema_diff_2']) & (self.data['close'] > 20)
        return self

    def buy_bottom(self):
        if len(self.data) < 25: # For HHV 25
            print(f"Group too small ({len(self.data)} rows) for buy_bottom; skipping.")
            self.data['BOTTOMBUY'] = False
            self.data['BOTTOMUPBUY'] = False
            self.data['POWERDOWNSELL'] = False
            self.data['POWERLINE'] = np.nan
            return self

        required_cols = ['low', 'high', 'close', 'open']
        missing_cols = [col for col in required_cols if col not in self.data.columns]
        if missing_cols:
            print(f"Missing columns {missing_cols}; skipping buy_bottom.")
            self.data['BOTTOMBUY'] = False
            self.data['BOTTOMUPBUY'] = False
            self.data['POWERDOWNSELL'] = False
            self.data['POWERLINE'] = np.nan
            return self

        self.data['VAR2'] = LLV(self.data['low'],10)
        self.data['VAR3'] = HHV(self.data['high'],25)
        
        # Avoid division by zero for POWERLINE
        denominator_powerline = self.data['VAR3'] - self.data['VAR2']
        # --- START MODIFICATION ---
        denominator_powerline = np.where(denominator_powerline == 0, 0.0001, denominator_powerline)
        # --- END MODIFICATION ---

        self.data['POWERLINE'] = np.round(EMA((self.data['close'] - self.data['VAR2']) / denominator_powerline * 4, 4),2)
        self.data['POWERDOWNSELL'] = (self.data['POWERLINE'].shift(1) > 1.75) & (self.data['POWERLINE'] < 1.75)
        
        if len(self.data) < 20: # For MA20
            print(f"Group too small ({len(self.data)} rows) for MA20 in buy_bottom; skipping some calculations.")
            self.data['BOTTOMBUY'] = False # Set to False if conditions cannot be fully met
            self.data['BOTTOMUPBUY'] = False
            return self

        self.data['MA20'] = self.data['close'].rolling(20).mean()
        
        # Avoid division by zero for WW
        denominator_ww = self.data['MA20']
        # --- START MODIFICATION ---
        denominator_ww = np.where(denominator_ww == 0, 0.0001, denominator_ww)
        # --- END MODIFICATION ---

        self.data['WW'] = (ABS((2*self.data['close'] + self.data['high'] + 2*self.data['low'])/4 - self.data['MA20']) / denominator_ww).clip(lower=0.01, upper=0.99)
        self.data['MM'] = DMA_torch(self.data['close'], self.data['WW'])
        self.data['通道'] = (1 - 7 / 100) * self.data['MM']
        self.data['VAR4'] = EMA(self.data['close'], 9)
        self.data['LDN'] = EMA(self.data['VAR4'] * 0.86,5)
        self.data['TYP'] = (self.data['high'] + self.data['low'] + self.data['close']) / 3
        
        # CCI calculation - MyTT does not have CCI, we will calculate it manually or add it if needed
        # Assuming you have a CCI calculation somewhere or add it here
        # For simplicity, let's add a dummy CCI if not present, but for real usage, implement CCI
        if 'CCI' not in self.data.columns:
            # Example CCI calculation (if not provided externally)
            typical_price = (self.data['high'] + self.data['low'] + self.data['close']) / 3
            ma_typical_price = typical_price.rolling(window=14).mean()
            mean_deviation = typical_price.rolling(window=14).apply(lambda x: np.mean(np.abs(x - np.mean(x))), raw=True)
            # Avoid division by zero
            mean_deviation_safe = mean_deviation.replace(0, np.nan).fillna(0.0001)
            self.data['CCI'] = (typical_price - ma_typical_price) / (0.015 * mean_deviation_safe)
            self.data['CCI'] = self.data['CCI'].replace([np.inf, -np.inf], np.nan).fillna(0) # Handle inf values

        ref_close = self.data['close'].shift(1)
        high_low = self.data['high'] - self.data['low']
        high_ref = (self.data['high'] - ref_close).abs()
        ref_low = (ref_close - self.data['low']).abs()
        tr = pd.concat([high_low, high_ref, ref_low], axis=1).max(axis=1)
        self.data['TTRR'] = tr.rolling(14).max()
        self.data['HD'] = self.data['high'] - REF(self.data['high'],1)
        self.data['LD'] = REF(self.data['low'],1) - self.data['low']  
        
        # Handle cases where TTRR might be 0, leading to division by zero in PDI/MDI
        dmp_val = IF((self.data['HD'] > 0) & (self.data['HD'] > self.data['LD']), self.data['HD'], 0)
        dmm_val = IF((self.data['LD'] > 0) & (self.data['LD'] > self.data['HD']), self.data['LD'], 0)

        self.data['DMP'] = EXPMEMA(dmp_val, 14)
        self.data['DMM'] = EXPMEMA(dmm_val, 14)

        ttrr_safe = self.data['TTRR']
        # --- START MODIFICATION ---
        ttrr_safe = np.where(ttrr_safe == 0, 0.0001, ttrr_safe)
        # --- END MODIFICATION ---
        self.data['PDI'] = self.data['DMP']*100/ttrr_safe
        self.data['MDI'] = self.data['DMM']*100/ttrr_safe
        
        denominator_adx = (self.data['MDI']+self.data['PDI'])
        # --- START MODIFICATION ---
        denominator_adx = np.where(denominator_adx == 0, 0.0001, denominator_adx)
        # --- END MODIFICATION ---

        self.data['ADX'] = EXPMEMA(ABS(self.data['MDI']-self.data['PDI'])/denominator_adx*100,6)

        condition1 = (self.data['close'] <= self.data['通道'])
        condition2 = (self.data['close'] < self.data['LDN'])
        condition3 = (self.data['CCI'] < -100)
        condition4 = (self.data['ADX'] >= 60)
        condition5 = (self.data['MDI'] >= self.data['PDI'])
        condition6 = (self.data['close'] > self.data['open'])
        satisfied_count = (
            condition1.astype(int) + 
            condition2.astype(int) + 
            condition3.astype(int) + 
            condition4.astype(int) + 
            condition5.astype(int) + 
            condition6.astype(int)
        )
        self.data['BOTTOMBUY'] = satisfied_count >= 4
        self.data['BOTTOMUPBUY'] = CROSS_TORCH(self.data['POWERLINE'], 0.5)
        return self

    def sell_on_top(self):
        """
        计算动力线及底部/关注信号，并结合放量+价格变化智能优化清仓和减仓条件
        """
        if len(self.data) < 5: # For rolling window 5
            print(f"Group too small ({len(self.data)} rows) for sell_on_top; skipping.")
            self.data['CLEANSELL'] = False
            self.data['STAGESELL'] = False
            self.data['ddx'] = np.nan
            self.data['POWERDOWNSELL'] = False # This is also set in buy_bottom, but just in case.
            return self

        required_cols = ['volume', 'close', 'open']
        missing_cols = [col for col in required_cols if col not in self.data.columns]
        if missing_cols:
            print(f"Missing columns {missing_cols}; skipping sell_on_top.")
            self.data['CLEANSELL'] = False
            self.data['STAGESELL'] = False
            self.data['ddx'] = np.nan
            return self

        large_order_threshold = 100 * 100  # 假设1手=100股
        # 计算大单净买入量（DDX）
        self.data['ddx'] = (
            (self.data['volume'] * (self.data['close'] - self.data['open']).apply(np.sign))
            .rolling(window=5)
            .sum()
        )
        # 计算五日均量和量比
        self.data['量5均'] = self.data['volume'].rolling(window=5).mean()
        
        # Avoid division by zero for 五日量比
        vol_5_mean_safe = self.data['量5均']
        # --- START MODIFICATION ---
        vol_5_mean_safe = np.where(vol_5_mean_safe == 0, 0.0001, vol_5_mean_safe)
        # --- END MODIFICATION ---
        self.data['五日量比'] = self.data['volume'] / vol_5_mean_safe
        
        # 计算跌幅
        self.data['跌幅'] = (self.data['close'] / self.data['close'].shift(1)) - 1

        # 计算DDX转负：今天<0 且 昨天>0
        self.data['DDX转负'] = (self.data['ddx'] < 0) & (self.data['ddx'].shift(1) > 0)

        # 判断是否放量：今日成交量 > 五日均量
        self.data['放量'] = self.data['volume'] > self.data['量5均']

        # 清仓线与减仓线
        清仓线 = 3.5
        减仓线 = 3.2

        # 清仓条件1：powerline 完全离开清仓线
        # Ensure 'POWERLINE' is calculated. It's done in buy_bottom.
        if 'POWERLINE' not in self.data.columns:
            print("POWERLINE not found, running buy_bottom to calculate it for sell_on_top.")
            self.buy_bottom() # Ensure POWERLINE is calculated
            if 'POWERLINE' not in self.data.columns: # If still not found, then data is too small for buy_bottom
                print("POWERLINE still not available after trying to calculate in sell_on_top. Skipping sell signals.")
                self.data['CLEANSELL'] = False
                self.data['STAGESELL'] = False
                return self

        self.data['清仓离开'] = (
            (self.data['POWERLINE'].shift(1) >= 清仓线) &
            (self.data['POWERLINE'] < 清仓线)
        )

        # 清仓条件2：POWERLINE 连续两天下降或持平
        self.data['POWERDOWN'] = self.data['POWERLINE'] <= self.data['POWERLINE'].shift(1)
        self.data['POWERDOWN2'] = (
            self.data['POWERDOWN'] & self.data['POWERDOWN'].shift(1)
        )

        # 最终清仓条件：
        self.data['CLEANSELL'] = (
            (self.data['清仓离开'].astype(int) +
             self.data['DDX转负'].astype(int) +
             self.data['放量'].astype(int) +
             (self.data['跌幅'] < 0).astype(int) +
             self.data['POWERDOWN2'].astype(int)) >= 4
        )

        # 减仓条件：powerline 触碰减仓线立刻回落
        self.data['STAGESELL'] = (
            (self.data['POWERLINE'].shift(1) > 减仓线) &
            (self.data['POWERLINE'] < 减仓线)
        )

        return self
    
    def buy_strong(self):
        if len(self.data) < 60: # For MA60
            print(f"Group too small ({len(self.data)} rows) for buy_strong; skipping.")
            self.data['STRONGBUY'] = False
            return self

        # Ensure WR6, K, D, CCI, RSI1, RSI2 are calculated
        # For simplicity, if not present, assume they are False or calculate them if possible.
        # It's better to calculate them explicitly or ensure they are provided in initial data.
        # Let's add dummy values or a placeholder if missing for now.
        for col in ['WR6', 'K', 'D', 'CCI', 'RSI1', 'RSI2']:
            if col not in self.data.columns:
                print(f"Warning: Column '{col}' not found for buy_strong. Setting to default/dummy values.")
                # You should ideally implement these calculations if they are fundamental for STRONGBUY
                # For this example, we'll set them to a value that is unlikely to trigger the signal
                if col in ['WR6', 'K', 'D']:
                    self.data[col] = 100 # High values to prevent WR6<5, K<50
                elif col == 'CCI':
                    self.data[col] = 0 # Not <-90 or >-100
                elif col in ['RSI1', 'RSI2']:
                    self.data[col] = 60 # Not <50

        self.data['buy_strong_d1'] = self.data['WR6'] < 5
        self.data['buy_strong_d2'] = CROSS_TORCH(self.data['K'], self.data['D']) & (self.data['K'] < 50)
        macd_dif, macd_dea, macd = MACD(self.data['close']) # MACD is from MyTT
        self.data['buy_strong_d3'] = CROSS_TORCH(macd_dif, macd_dea)
        self.data['buy_strong_d4'] = (self.data['CCI'] < -90) & (self.data['CCI'] > -100)
        self.data['buy_strong_d5'] = CROSS_TORCH(self.data['RSI1'], self.data['RSI2']) & (self.data['RSI1'] < 50)
        self.data['buy_strong_aaa'] = self.data['close'] > self.data['open']
        self.data['buy_strong_bbb'] = self.data['buy_strong_aaa'] & \
                                     (self.data['close'] > MA(self.data['close'], 20)) & \
                                     (self.data['close'] > MA(self.data['close'], 40)) & \
                                     (self.data['close'] > MA(self.data['close'], 60))
        self.data['buy_strong_ccc'] = self.data['buy_strong_bbb'] & \
                                     (self.data['open'] < MA(self.data['close'], 40)) & \
                                     (self.data['open'] < MA(self.data['close'], 60))
        self.data['buy_strong_xg1'] = self.data['buy_strong_ccc'] & \
                                     ((self.data['close'] - self.data['open']) > 0.0618 * self.data['close'])
        self.data['buy_strong_aa1'] = self.data['buy_strong_d1'] & self.data['buy_strong_d2']
        self.data['buy_strong_aa2'] = self.data['buy_strong_d2'] & self.data['buy_strong_d3']
        self.data['buy_strong_aa3'] = self.data['buy_strong_d1'] & self.data['buy_strong_d3']
        self.data['buy_strong_aa4'] = self.data['buy_strong_d4'] & self.data['buy_strong_d2']
        self.data['buy_strong_aa5'] = self.data['buy_strong_d5'] & self.data['buy_strong_d3']
        self.data['buy_strong_aa6'] = self.data['buy_strong_d4'] & self.data['buy_strong_d5']
        self.data['buy_strong_xg2'] = (
            self.data['buy_strong_aa1'] | 
            self.data['buy_strong_aa2'] | 
            self.data['buy_strong_aa3'] |
            self.data['buy_strong_aa4'] |
            self.data['buy_strong_aa5'] |
            self.data['buy_strong_aa6']
        )
        self.data['STRONGBUY'] = self.data['buy_strong_xg2']
        return self

    def buy_sell_trend(self):
        N = 20
        M = 32
        P1 = 80
        P2 = 100

        if len(self.data) < max(N, M):
            print(f"Group too small ({len(self.data)} rows) for buy_sell_trend; skipping.")
            self.data['TRENDBUY'] = False
            self.data['TRENDSELL'] = False
            return self

        if all(col in self.data.columns for col in ['close', 'high', 'open', 'low', 'volume']):
            self.data['VAR1_BSET'] = (self.data['close'] + self.data['high'] + self.data['open'] + self.data['low']) / 4
            self.data['SELL_LINE'] = XMA(self.data['VAR1_BSET'], N) * (1 + P1 / 1000)
            self.data['BUY_LINE'] = XMA(self.data['VAR1_BSET'], M) * (1 - P2 / 1000)
            self.data['TRENDBUY'] = (
                ((self.data['open'] < self.data['BUY_LINE']) | (self.data['low'] < self.data['BUY_LINE'])) &
                (self.data['close'] > self.data['BUY_LINE'])
            )
            sell_break_today = self.data['high'] > self.data['SELL_LINE']
            volume_expand_today = self.data['volume'] >= REF(self.data['volume'], 1) * 2
            self.data['TRENDSELL'] = sell_break_today & volume_expand_today
        return self
    
    def buy_dragon_catching(self):
        """计算擒龙追涨买点信号"""
        if len(self.data) < 2:
            print(f"Group too small ({len(self.data)} rows) for buy_dragon_catching; skipping.")
            self.data['DRAGONBUY'] = False
            return self

        # Ensure REF shifts the series correctly
        def REF(series, periods):
            return series.shift(periods)
    
        # Calculate VAR20: close > previous close AND close < previous close * 1.095
        self.data['VAR20'] = (self.data['close'] > REF(self.data['close'], 1)) & (self.data['close'] < REF(self.data['close'], 1) * 1.095)
    
        # Calculate DRAGONBUY signal
        self.data['DRAGONBUY'] = (
            (self.data['volume'] >= REF(self.data['volume'], 1) * 3) &
            (self.data['close'] < 10) &
            self.data['VAR20'] &
            (self.data['close'] != self.data['high']) &
            (self.data['close'] > self.data['open'])
        )
        return self.data['DRAGONBUY']

    def fishing_season_indicator(self):
        """
        通达信【捕捞季节】指标
        """
        required_cols = ['close', 'high', 'low', 'amount', 'volume']
        missing_cols = [col for col in required_cols if col not in self.data.columns]
        if missing_cols:
            print(f"Missing columns {missing_cols} for fishing_season_indicator; skipping.")
            # Set default values for all generated columns
            self.data['WY1001'] = np.nan
            self.data['WY1002'] = np.nan
            self.data['WY1003'] = np.nan
            self.data['WY1004'] = np.nan
            self.data['XYS0'] = np.nan
            self.data['PJGJ'] = np.nan
            self.data['SSRYDJX'] = np.nan
            self.data['SSRCJL'] = np.nan
            self.data['SSRCJE'] = np.nan
            self.data['SSRCBJX'] = np.nan
            self.data['CYS13'] = np.nan
            self.data['XYSHSL'] = np.nan
            self.data['XYS0_STICKLINE_RED'] = False
            self.data['XYS0_STICKLINE_GREEN'] = False
            self.data['FISHING_SEASON_GREEN'] = False
            self.data['FISHING_SEASON_YELLOW'] = False
            self.data['FISHING_SEASON_CYAN'] = False
            self.data['FISHING_SEASON_BLUE'] = False
            self.data['XYS1'] = np.nan
            self.data['XYS2'] = np.nan
            return self

        # WY1001:=(2CLOSE+HIGH+LOW)/3;
        self.data['WY1001'] = (2 * self.data['close'] + self.data['high'] + self.data['low']) / 3

        # WY1002:=EMA(self.data['WY1001'],3);
        self.data['WY1002'] = EMA(self.data['WY1001'], 3)
        # WY1003:=EMA(self.data['WY1002'],3);
        self.data['WY1003'] = EMA(self.data['WY1002'], 3)
        # WY1004:=EMA(self.data['WY1003'],3);
        self.data['WY1004'] = EMA(self.data['WY1003'], 3)

        # XYS0:=(WY1004-REF(WY1004,1))/REF(WY1004,1)*100;
        ref_wy1004_1 = REF(self.data['WY1004'], 1)
        # Handle division by zero for XYS0
        # --- START MODIFICATION ---
        ref_wy1004_1_safe = np.where(ref_wy1004_1 == 0, 0.0001, ref_wy1004_1)
        # --- END MODIFICATION ---
        self.data['XYS0'] = (self.data['WY1004'] - ref_wy1004_1) / ref_wy1004_1_safe * 100

        # STICKLINE(XYS0>=0,XYS0,0,2,0),COLORRED;
        # STICKLINE(XYS0<0,XYS0,0,2,0),COLORGREEN;
        # For STICKLINE, we'll create boolean flags or actual values to represent the drawing.
        # Here, let's create boolean flags indicating the condition.
        self.data['XYS0_STICKLINE_RED'] = self.data['XYS0'] >= 0
        self.data['XYS0_STICKLINE_GREEN'] = self.data['XYS0'] < 0

        # PJGJ:=AMOUNT/VOL/100; (Note: VOL in MyTT is volume, AMOUNT is amount, typically in yuan, VOL in shares/lots)
        # Assuming 'vol' is in shares, and 'amount' is in yuan.
        # If 'vol' is lots, then 'amount'/'vol' is per lot. Need to clarify unit of 'vol'.
        # If Volume is in shares, then PJGJ = Amount / Volume.
        # Given `PJGJ:=AMOUNT/VOL/100;`, it suggests VOL is in shares and AMOUNT is in yuan, and then divided by 100 for price per share.
        
        # Handle division by zero for PJGJ
        vol_safe = self.data['volume']
        # --- START MODIFICATION ---
        vol_safe = np.where(vol_safe == 0, 0.0001, vol_safe)
        # --- END MODIFICATION ---
        self.data['PJGJ'] = self.data['amount'] / vol_safe / 100 

        # SSRYDJX:=MA(PJGJ,13);
        self.data['SSRYDJX'] = MA(self.data['PJGJ'], 13)
        # SSRCJL:=EMA(VOL,13);
        self.data['SSRCJL'] = EMA(self.data['volume'], 13)
        # SSRCJE:=EMA(AMOUNT,13);
        self.data['SSRCJE'] = EMA(self.data['amount'], 13)
        
        # SSRCBJX:=SSRCJE/SSRCJL/100;
        # Handle division by zero for SSRCJJX
        ssrcjl_safe = self.data['SSRCJL']
        # --- START MODIFICATION ---
        ssrcjl_safe = np.where(ssrcjl_safe == 0, 0.0001, ssrcjl_safe)
        # --- END MODIFICATION ---
        self.data['SSRCBJX'] = self.data['SSRCJE'] / ssrcjl_safe / 100

        # CYS13:=(CLOSE-SSRCBJX)/SSRCBJX*100;
        # Handle division by zero for CYS13
        ssrcbjx_safe = self.data['SSRCBJX']
        # --- START MODIFICATION ---
        ssrcbjx_safe = np.where(ssrcbjx_safe == 0, 0.0001, ssrcbjx_safe)
        # --- END MODIFICATION ---
        self.data['CYS13'] = (self.data['close'] - self.data['SSRCBJX']) / ssrcbjx_safe * 100

        # XYSHSL:=EMA(VOL/CAPITAL*100,13);
        # Assuming 'CAPITAL' (total shares/market cap) is available in self.data or needs to be provided.
        # If CAPITAL is not available, this calculation will fail.
        # For demonstration, let's assume 'CAPITAL' is a column. If not, it needs to be added to the input data.
        # If CAPITAL is not consistently available or varies greatly, this part might need adjustment.
        # For this example, if 'CAPITAL' is missing, we'll print a warning and skip relevant calculations.
        if 'capital' not in self.data.columns:
            print("Warning: 'capital' column not found for XYSHSL calculation. Skipping this part of fishing_season_indicator.")
            self.data['XYSHSL'] = np.nan
        else:
            # Handle division by zero for XYSHSL
            capital_safe = self.data['capital']
            # --- START MODIFICATION ---
            capital_safe = np.where(capital_safe == 0, 0.0001, capital_safe)
            # --- END MODIFICATION ---
            self.data['XYSHSL'] = EMA(self.data['volume'] / capital_safe * 100, 13)

        # STICKLINE conditions (converted to boolean flags for analysis)
        self.data['FISHING_SEASON_GREEN'] = (self.data['XYSHSL'] > 6.1) & (self.data['CYS13'] > 5)
        self.data['FISHING_SEASON_YELLOW'] = (self.data['XYSHSL'] > 3.8) & (self.data['CYS13'] > 5)
        self.data['FISHING_SEASON_CYAN'] = (self.data['XYSHSL'] > 2.1) & (self.data['CYS13'] > 5)
        self.data['FISHING_SEASON_BLUE'] = (self.data['XYSHSL'] > 1.8) & (self.data['CYS13'] > 5)

        # XYS1:MA(XYS0,1),COLOREA91B5,LINETHICK2;
        self.data['XYS1'] = MA(self.data['XYS0'], 1) # MA(X,1) is essentially X
        # XYS2:MA(XYS0,2),COLORYELLOW,LINETHICK1;
        self.data['XYS2'] = MA(self.data['XYS0'], 2)
        
        return self


    def calculate_all_signals(self):
            self.buy_bottom()
            self.buy_potential()
            self.three_stars_bottom()
            self.buy_entry_signal()
            #self.buy_sell_trend() # Commented out as per original file
            self.sell_on_top()
            self.buy_dragon_catching()
            self.fishing_season_indicator() # Add the new indicator calculation here
            return self.data
